/*!
  @file dense_stereo_graphcut.cpp
  @copyright 2013 Kubota Lab. All rights resereved.
*/

#include <ros/ros.h>

#include <limits>

#include <cv_bridge/cv_bridge.h>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

#include "stereo/dense_stereo_graphcut.h"


namespace stereo 
{

void DenseStereoGraphcut::stereoImageCB(
    const sensor_msgs::ImageConstPtr &imgL, const sensor_msgs::CameraInfoConstPtr &infoL, 
    const sensor_msgs::ImageConstPtr &imgR, const sensor_msgs::CameraInfoConstPtr &infoR)
{
  {
    boost::mutex::scoped_lock lock(mutex_);
    if (!trigger_) 
    {
      trigger_ = true;
      return;
    }
  }

  if (disp_pub_.getNumSubscribers() > 0)
  {
    {
      boost::mutex::scoped_lock lock(mutex_);
      if (processing_) return;
      processing_ = true;
    }

    model_.fromCameraInfo(infoL, infoR);
    ROS_INFO_STREAM(model_.left().fx());
    ROS_INFO_STREAM(model_.baseline());
    if (dispmax_ < 0 || dispmin_ < 0)
    {
      dispmax_ = scale_ * model_.getDisparity(distmin_) + 1;
      dispmin_ = scale_ * model_.getDisparity(distmax_);
      ROS_INFO_STREAM("Disparity: " << dispmin_ << "~" << dispmax_);
    }

    cv::Mat _imL = cv_bridge::toCvShare(imgL)->image;
    cv::Mat _imR = cv_bridge::toCvShare(imgR)->image;

    /* resize image if necessary */
    int width  = 1.0 * _imL.cols * scale_;
    int height = 1.0 * _imL.rows * scale_;
    cv::Mat imL(height, width, CV_8U);
    cv::Mat imR(height, width, CV_8U);
    cv::resize(_imL, imL, imL.size());
    cv::resize(_imR, imR, imR.size());
    //cv::Mat data(height, width, CV_32F, dispmax_/2);
    cv::Mat data(height, width, CV_32F, -1);

    process(imL, imR, data);

    /* publish disparity */
    {
      stereo_msgs::DisparityImage::Ptr disparity(new stereo_msgs::DisparityImage);
      disparity->header = imgL->header;
      cv_bridge::CvImage(disparity->header, sensor_msgs::image_encodings::TYPE_32FC1, data).toImageMsg(disparity->image);
      disparity->min_disparity = 0;
      disparity->max_disparity = dispmax_;
      disparity->f = scale_ * 0.5 * (model_.left().fx() + model_.left().fy());
      disparity->T = model_.baseline();
      disparity->delta_d = 1.0;  // tmp
      disp_pub_.publish(disparity);
    }


    {
      boost::mutex::scoped_lock lock(mutex_);
      processing_ = false;
      trigger_ = false;
    }
  }
}

void DenseStereoGraphcut::process(const cv::Mat &imgL, const cv::Mat &imgR, cv::Mat &data)
{
  int width  = imgL.cols;
  int height = imgL.rows;

  cv::Mat imL = imgL;
  cv::Mat imR = imgR;
  if (imL.channels() == 3)
  {
    cv::Mat imL_cp = imL;
    cv::Mat imR_cp = imR;
    cv::cvtColor(imL_cp, imL, CV_BGR2GRAY);
    cv::cvtColor(imL_cp, imL, CV_BGR2GRAY);
  }

  if (do_sobel_)
  {
    if (imL.channels() == 3)
    {
      cv::Mat imL_cp = imL;
      cv::Mat imR_cp = imR;
      cv::cvtColor(imL_cp, imL, CV_BGR2GRAY);
      cv::cvtColor(imL_cp, imL, CV_BGR2GRAY);
    }
    performSobel(imL, imL);
    performSobel(imR, imR);
  }

  fprintf(stdout, "Computing disparities for %dx%d images....\n", width, height);
  clock_t start = clock();
  {
    size_t num_node = width * height;
    size_t num_edge = (width - 1) * height + width * (height - 1);
    size_t max_num_node = num_node + num_edge;
    size_t max_num_edge = (width+(width-1)-1)*height+width*(height+(height-1)-1);
    boost::shared_ptr<GraphType> g(new GraphType(max_num_node, max_num_edge));

    double E = 10e10;  // initial value of the cost function
    size_t LOOPMAX = 10e5;
    for (size_t loopidx = 0; loopidx < LOOPMAX; ++loopidx)
    {
      fprintf(stdout, "Loop %lu\n", loopidx);
      bool success = false;
      for (float a = 0; a < dispmax_; ++a)
      {
        g->add_node(num_node);

        /* data term */
        for (int i = 0; i < height; ++i)
          for (int j = 0; j < width; ++j)
          {
            double d_p = sad(j, i, imL, imR, data.at<float>(i, j));
            double d_a = sad(j, i, imL, imR, a);
            g->add_tweights(i * width + j, d_p, d_a);
          }

        /* smoothness term */
        // vertical pixel link
        size_t added_node = 0;
        for (int i = 0; i < height - 1; ++i)
          for (int j = 0; j < width; ++j)
          {
            double pp = paramPotts(
                (imL.channels() == 3)? imL.at<cv::Vec3b>(i, j)[0]:   imL.at<uchar>(i, j),
                (imL.channels() == 3)? imL.at<cv::Vec3b>(i+1, j)[0]: imL.at<uchar>(i+1, j));

            if (data.at<float>(i, j) == data.at<float>(i+1, j))
            {
              double v = wdiff(a, data.at<float>(i, j), pp);
              g->add_edge(i*width+j, (i+1)*width+j, v, v);
            }
            else
            {
              double vs = wdiff(data.at<float>(i, j), data.at<float>(i+1, j), pp);
              double vt = 0;  // =diff(a, a)
              g->add_node();
              int new_node = num_node + added_node;
              ++added_node;
              g->add_tweights(new_node, vs, vt);

              double v0 = wdiff(a, data.at<float>(i, j), pp);
              double v1 = wdiff(a, data.at<float>(i+1, j), pp);
              g->add_edge(i*width+j, new_node, v0, v0);
              g->add_edge((i+1)*width+j, new_node, v1, v1);
            }
          }
        // horizontal pixel link
        for (int i = 0; i < height; ++i)
          for (int j = 0; j < width - 1; ++j)
          {
            double pp = paramPotts(
                (imL.channels() == 3)? imL.at<cv::Vec3b>(i, j)[0]:   imL.at<uchar>(i, j),
                (imL.channels() == 3)? imL.at<cv::Vec3b>(i, j+1)[0]: imL.at<uchar>(i, j+1));

            if (data.at<float>(i, j) == data.at<float>(i, j+1))
            {
              double v = wdiff(a, data.at<float>(i, j), pp);
              g->add_edge(i*width+j, i*width+j+1, v, v);
            }
            else
            {
              double vs = wdiff(data.at<float>(i, j), data.at<float>(i, j+1), pp);
              double vt = 0;  // =diff(a, a)
              g->add_node();
              int new_node = num_node + (added_node++);
              g->add_tweights(new_node, vs, vt);

              double v0 = wdiff(a, data.at<float>(i, j), pp);
              double v1 = wdiff(a, data.at<float>(i, j+1), pp);
              g->add_edge(i, new_node, v0, v0);
              g->add_edge(i+1, new_node, v1, v1);
            }
          }

        /* max flow, min cut */
        double E_new = g->maxflow();
        //fprintf(stdout, "Alpha %d  Node %d  Arc %d  Flow %.2f\n", a, g->get_node_num(), g->get_arc_num(), E_new);
        if (E_new < E)
        {
          E = E_new;
          success = true;
          for (int i = 0; i < height; ++i)
            for (int j = 0; j < width; ++j)
            {
              if (g->what_segment(i*width+j) == GraphType::SOURCE) 
                data.at<float>(i, j) = a;
            }
        }

        /* finishing */
        g->reset();
      }
      if (!success) break;
    }

    /* subpixel */
    if (1)
    {
      for (int i = 0; i < height; ++i)
        for (int j = 0; j < width; ++j)
        {
          if (data.at<float>(i, j) > 0)
          {
            ad(j, i, imL, imR, data.at<float>(i, j), true);
          }
        }

    }
  }
  clock_t end = clock();
  fprintf(stdout, "elapsed time: %.2f [s]\n", ((double)(end-start))/CLOCKS_PER_SEC);
}

void DenseStereoGraphcut::performSobel(const cv::Mat &src, cv::Mat &dst)
{
  cv::Mat src_g(src.size(), CV_8U);
  cv::Mat grad_x, grad_y;
  cv::Mat abs_grad_x, abs_grad_y;
  if (src.channels() == 3)
    cvtColor(src, src_g, CV_BGR2GRAY);
  else
    src_g = src;

  Sobel(src_g, grad_x, CV_16S, 1, 0, 3, 1, 0, cv::BORDER_DEFAULT);
  convertScaleAbs(grad_x, abs_grad_x);
  Sobel(src_g, grad_y, CV_16S, 0, 1, 3, 1, 0, cv::BORDER_DEFAULT);
  convertScaleAbs(grad_y, abs_grad_y);

  if (!dst.data) dst.create(src.size(), CV_8U);
  addWeighted(abs_grad_x, 0.5, abs_grad_y, 0.5, 0, dst);
}

double DenseStereoGraphcut::sad(const int x, const int y, const cv::Mat &imL, const cv::Mat &imR, const int dsp)
{
  int count = 0;
  int sum = 0;
  int winsize = winsize_;
  for (int i = - winsize/2; i <= winsize/2; ++i)
    for (int j = -winsize/2; j <= winsize/2; ++j)
      if (isBetween(y+i, 0, imL.rows) 
          && isBetween(x+j, 0, imL.cols) 
          && isBetween(x+j-dsp, 0, imL.cols))
      {
        ++count;
        if (imL.channels() == 3)
          sum += abs(imL.at<cv::Vec3b>(y+i, x+j)[0] - imR.at<cv::Vec3b>(y+i, x+j-dsp)[0])
               + abs(imL.at<cv::Vec3b>(y+i, x+j)[1] - imR.at<cv::Vec3b>(y+i, x+j-dsp)[1])
               + abs(imL.at<cv::Vec3b>(y+i, x+j)[2] - imR.at<cv::Vec3b>(y+i, x+j-dsp)[2]);
        else
          sum += abs(imL.at<uchar>(y+i, x+j) - imR.at<uchar>(y+i, x+j-dsp));
      }
  return count? 1.0 * sum / (count * imL.channels()): std::numeric_limits<double>::max();
}

double DenseStereoGraphcut::ad(const int x, const int y, const cv::Mat &imL, const cv::Mat &imR, float &d, bool rewrite)
{
  double Cfwd = std::numeric_limits<double>::max();
  double Crev = std::numeric_limits<double>::max();
  double min_d = 0.0;
  for (double dtic = - 0.5; dtic <= 0.5; dtic += 0.1)
  {
    double prev = Cfwd;
    if (imL.channels() == 3)
      Cfwd = std::min(Cfwd, fabsN(interpolateN<double, 3>(imL, x, y), interpolateN<double, 3>(imR, x-(d+dtic), y)));
    else
      Cfwd = std::min(Cfwd, fabs(interpolate(imL, x, y) - interpolate(imR, x-(d+dtic), y)));
    if (prev != Cfwd) min_d = d + dtic;
  }
  double min_xtic = 0.0;
  for (double xtic = - 0.5; xtic <= 0.5; xtic += 0.1)
  {
    double prev = Crev;
    if (imL.channels() == 3)
      Crev= std::min(Crev, fabsN(interpolateN<double, 3>(imL, x+xtic, y), interpolateN<double, 3>(imR, x-min_d, y)));
    else
      Crev= std::min(Crev, fabs(interpolate(imL, x+xtic, y) - interpolate(imR, x-min_d, y)));
    if (prev != Crev) min_xtic = xtic;
  }

  if (rewrite) d = min_d + min_xtic;
  return std::min(std::min(Cfwd, Crev), 20.0);
}

}  // namespace stereo
